import base64
import hashlib
import random
from binascii import a2b_hex, b2a_hex

import rsa

from common.utils.security import PyRSA
from common.api import http_result


def request_pk(step, pk=None, session=None):
    if step == '1':
        if 'pub_key' in session:
            # pub_key 的base64
            pub_key = session.get('pub_key')
            return http_result.result(1, http_result.RESULT_SUCCESS, {'key': pub_key})

        keys = PyRSA().gen_keys()
        # print(keys[1])
        pri_k = str(b2a_hex(keys[3]))[2: -1]
        session['pri_key'] = pri_k

        k2 = str(keys[2])[2:-1]
        # print('公钥', k2)

        pub_key = base64.b64encode(k2.encode()).decode()
        session['pub_key'] = pub_key

        return http_result.result(1, http_result.RESULT_SUCCESS, {'key': pub_key})

    elif step == '2':
        # 客户端传来的public key
        # print('pk: %s'%pk)
        # print('key', key)
        pk = base64.b64decode(pk)
        # 这个pk 为加密过的公钥  使用第一步中服务器传给客户端的公钥进行的加密
        # 解密这个公钥
        pri_key_str = session.get('pri_key')
        if not pri_key_str:
            return http_result.result(0, 'there is no private key')

        pri_key = PyRSA.parse_private_key(a2b_hex(pri_key_str))
        # print(pri_key)
        # 解密 :  客户端发过来的是用服务器发给客户端的公钥加密后的密文  需要解密得到客户端的公钥(明文)
        pub_key_str = rsa.decrypt(pk, pri_key)
        # print(pub_key_str)

        # 由公钥明文 -> 公钥对象
        pub_key = PyRSA.parse_public_key(pub_key_str)
        # print(pub_key)

        # 生成key
        key = gen_key()
        # print('aes key is: ', key)

        # 使用客户端的公钥加密
        key0 = rsa.encrypt(bytes(key, encoding='utf-8'), pub_key)
        key0 = base64.b64encode(key0).decode()
        # print(key0)

        session['aes_key'] = key
        result = http_result.result(1, http_result.RESULT_SUCCESS, {'key': key0})
        return result

    else:
        return http_result.result(0, 'step is illegal')


# 生成随机密码 32位
# 随机数的md5
def gen_key():
    rand = random.random()
    # print(rand)
    md = hashlib.md5()
    md.update(str(rand).encode())
    key = md.hexdigest()
    # print(key, type(key))
    return key
